import argparse

import torch

from config import read_config
from trainer import Trainer
from utils import seed_it, Plot


def main(path):
    config = read_config(path)
    seed_it(config['seed'])
    if config['plot']:
        plot = Plot(config['model'], config['dataset']['type'])
        plot.plot(step=config['plot_step'])
    else:
        trainer = Trainer(
            n_clients=config['clients_number'],
            epoch=config['epoch'],
            optimizer=config['optimizer']['type'],
            weight_decay=float(config['optimizer']['weight_decay']),
            momentum=float(config['optimizer']['momentum']),
            lr=config['optimizer']['learning_rate'],
            scheduler=config['scheduler']['type'],
            step_size=config['scheduler']['step_size'],
            gamma=config['scheduler']['gamma'],
            dataset=config['dataset']['type'],
            path=config['dataset']['path'],
            alpha=config['dataset']['alpha'],
            model=config['model'],
            batch_size=config['batch_size'],
            compress_policy=config['compression']['type'],
            ratio=config['compression']['ratio'],
            bit=config['compression']['bit'],
            client_device=torch.device(config['device']['client']),
            server_device=torch.device(config['device']['server']),
        )
        trainer.train_loop()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', default="./config.yml", type=str, help='config file')
    arg = parser.parse_args()
    main(arg.config)
